#!/bin/bash

set -o pipefail

# shellcheck disable=SC1091
. /usr/libexec/os-helpers-logging

# Help function
update_balena_supervisor_help() {
    cat << EOF
Wrapper to run supervisor agent updates on balena distributions.
$0 <OPTION>

Options:
  -h, --help
        Display this help and exit.

  -i <SUPERVISOR IMAGE>, --supervisor-image <SUPERVISOR IMAGE>
        Set supervisor image to update to. This can be an image path like:
          registry2.balena-cloud.com/v2/8f3d658a373039a9c27b8f1811b0261d
        or a friendly URL like:
          bh.cr/<org>/<fleet>/<commit|semver>

  -n, --no-start-stop-supervisor
        Do not start/stop the supervisor.
EOF
}

START_STOP_SUPERVISOR=1
# Parse arguments
while [ $# -gt 0 ]; do
    arg="$1"

    case $arg in
        -h|--help)
            update_balena_supervisor_help
            exit 0
            ;;
        -i|--supervisor-image)
            if [ -z "$2" ]; then
                fail "\"$1\" argument needs a value."
            fi
            # Matches <[sub.]domain.tld>/<fleet|translation>/<digest|semver>
            # For example:
            # bh.cr/balena/supervisor/0.0.0
            # registry2.balena-cloud.com/v2/8f3d658a373039a9c27b8f1811b0261d
            _regex='^([A-Za-z0-9\-]{1,63}\.)+[A-Za-z]{2,6}\/[a-zA-Z0-9_\.\-\/]+|v[0-9]+\/[0-9a-zA-Z_\-]+$'
            if [[ "$2" =~ ${_regex} ]]; then
                  TARGET_SUPERVISOR_IMAGE=$2
            else
                fail "\"$2\" invalid value - should be of the form <[sub.]domain.tld>/<fleet|translation/<digest|semver>."
            fi
            shift
            ;;

        -n|--no-start-stop-supervisor)
            START_STOP_SUPERVISOR=0
            shift
            ;;
        *)
            error "Unrecognized option $1."
            ;;
    esac
    shift
done

# Don't source before parsing args - balena-config-vars parses args too
# shellcheck disable=SC1091
. /usr/sbin/balena-config-vars

# A temporary file used until next reboot
UPDATECONF=/tmp/update-supervisor.conf
# The supervisor state configuration directory
SVCONFIGDIR=$(basename "$(find /etc -type d -name "*supervisor")")

# If the user api key exists we use it instead of the deviceApiKey as it means we haven't done the key exchange yet
_device_api_key=${PROVISIONING_API_KEY:-$DEVICE_API_KEY}

error_handler() {
    # shellcheck disable=SC2181
    [ $? -eq 0 ] && exit 0

    error "${1}"
    # If docker pull fails, start the old supervisor again and exit
    rm -rf $UPDATECONF
    if [ "${START_STOP_SUPERVISOR}" -eq 1 ]; then
        systemctl start balena-supervisor
    fi
    exit 1
}

trap error_handler EXIT

# Fetch target state from API
#
# Inputs:
# $1: Device UUID
# $2: Balena API environment
# $3: Balena API token
#
# Outputs:
# On success (0), space separated supervisor version and image name, or "null null" if there is no match
# On failure (1), no output
#
os_helpers_fetch_target_state() {
    _device_uuid="$1"
    _api_env="$2"
    _token="$3"

    # shellcheck disable=SC1091
    . /usr/libexec/os-helpers-api

    if _response=$(api_get_request "${_api_env}/v6/supervisor_release?\$select=supervisor_version,image_name&\$filter=should_manage__device/any(d:d/uuid%20eq%20'${_device_uuid}')" "${_token}"); then
        if _target_state_json=$(echo "${_response}" | jq -e -r '.d[0].supervisor_version,.d[0].image_name'); then
            echo "${_target_state_json}" | tr "\n" " "
            return 0
        fi
    fi
    return 1
}


# Detect containers engine
if which docker > /dev/null 2>&1; then
    DOCKER=docker
elif which rce > /dev/null 2>&1; then
    DOCKER=rce
elif which balena > /dev/null 2>&1; then
    DOCKER=balena
else
    error_handler "no container engine detected"
fi

# Get target supervisor details from API.
# The script will exit if curl does not get a valid response.
# Getting data separately before reading it fixes error handling.
info "Getting image name and version..."
if [ -n "$API_ENDPOINT" ] && [ -n "${UUID}" ] && [ -n "$_device_api_key" ]; then
    if _target_state=$(os_helpers_fetch_target_state "${UUID}" "${API_ENDPOINT}" "${_device_api_key}"); then
        if [ -n "${_target_state}" ]; then
            read -r version image_name <<< "${_target_state}"
            if [ -z "$version" ] || [ -z "$image_name" ]; then
                error_handler "unexpected API data"
            fi
        fi
    fi
fi

# shellcheck disable=SC1090
. "/etc/${SVCONFIGDIR}/supervisor.conf"

# If no API version is set, use preloaded values
if [ -z "$version" ] || [ -z "$image_name" ] || [ "$version" = "null" ] || [ "$image_name" = "null" ]; then
    info "No supervisor configuration found from API."
    if [ -n "${TARGET_SUPERVISOR_IMAGE}" ]; then
        info "Using command line image argument $TARGET_SUPERVISOR_IMAGE."
        image_name="${TARGET_SUPERVISOR_IMAGE}"
        version=$(basename "${image_name}")
    else
        info "Using preloaded values."
        if [ -z "$SUPERVISOR_VERSION" ]; then
            error_handler "no preloaded version found"
        fi
        if [ -z "$SUPERVISOR_IMAGE" ]; then
            error_handler "no preloaded image found."
        fi
        info "Set based on preloaded values image=$SUPERVISOR_IMAGE and version=$SUPERVISOR_VERSION."
        image_name="$SUPERVISOR_IMAGE"
        version="$SUPERVISOR_VERSION"
    fi
fi

setSupervisorConfig() {
    # Store the tagged image string so balena-supervisor.service can pick it up
    sed -e "s|SUPERVISOR_IMAGE=.*|SUPERVISOR_IMAGE=$image_name|" -e "s|SUPERVISOR_TAG|SUPERVISOR_VERSION|" -e "s|SUPERVISOR_VERSION=.*|SUPERVISOR_VERSION=$version|" "/etc/${SVCONFIGDIR}/supervisor.conf" > $UPDATECONF
}

modifySupervisorConfig() {
    sed -i -e "s|SUPERVISOR_IMAGE=.*|SUPERVISOR_IMAGE=$image_name|" -e "s|SUPERVISOR_TAG|SUPERVISOR_VERSION|" -e "s|SUPERVISOR_VERSION=.*|SUPERVISOR_VERSION=$version|" "/etc/${SVCONFIGDIR}/supervisor.conf"
    sync -f /mnt/state
}

info "Getting image id..."
imageid=$($DOCKER inspect -f '{{.Id}}' "$image_name") || imageid=""

if [ -n "$imageid" ]; then
    info "Supervisor $image_name at version $version already downloaded."
    # Since target Supervisor version is downloaded ensure the device is configured to use it
    setSupervisorConfig
    # Check if it's already running
    if [ -z "$(balena container ls --filter=ancestor="${imageid}" --format "{{.ID}}")" ]; then
        if [ "${START_STOP_SUPERVISOR}" -eq 1 ]; then
            systemctl restart balena-supervisor
        else
            warn "Target supervisor downloaded but not running and no supervisor restart was specified."
        fi
    fi
    # Make sure supervisor.conf has been updated
    if ! grep -q "SUPERVISOR_IMAGE=$image_name" "/etc/${SVCONFIGDIR}/supervisor.conf"; then
        warn "Setting supervisor.conf to use $image_name at $version."
        modifySupervisorConfig
    fi
    exit 0
fi

if [ "${START_STOP_SUPERVISOR}" -eq 1 ]; then
    # Try to stop old supervisor to prevent it deleting the intermediate images while downloading the new one
    info "Stop supervisor..."
    systemctl stop balena-supervisor
fi

# Pull target version.
info "Pulling supervisor $image_name at version $version..."
if _out=$($DOCKER pull "$image_name"); then
    info "${_out}"
    image_id=$(balena images --filter=reference="${image_name%@*}" --format "{{.ID}}")
    $DOCKER tag "${image_id}" "balena_supervisor:${version}"
else
    error "${_out}"
    error_handler "supervisor pull failed"
fi

# Ensure device has newly pulled image as the version to run
setSupervisorConfig

# Run supervisor with the device-type-specific options.
# We give a specific name to the container to guarantee only one running.
if [ "${START_STOP_SUPERVISOR}" -eq 1 ]; then
    info "Start supervisor..."
    systemctl start balena-supervisor
fi

# Make sure supervisor.conf has been updated
modifySupervisorConfig
